# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import json
import logging
import enum

from tornado import gen
from tornado.ioloop import IOLoop
from tornado.websocket import websocket_connect, WebSocketClosedError
from tornado.concurrent import Future

from coin.base.timestamp import get_timestamp
from coin.exchange.base.rest.task_scheduler import Scheduler
from coin.exchange.kraken.kr_rest.native_private_client import KrakenAsyncAuth

KRAKEN_FUTURES_WS_API_URL = 'wss://ws-auth.kraken.com'


def pretty_json(data):
  return json.dumps(data, indent=2)


class ChannelType(enum.Enum):
  OPEN_ORDERS = 'openOrders'
  OWN_TRADES = 'ownTrades'


class EventType(enum.Enum):
  PING = 'ping'
  PONG = 'pong'
  HEARTBEAT = 'heartbeat'
  SYSTEM_STATUS = 'systemStatus'
  SUBSCRIPTION_STATUS = 'subscriptionStatus'
  SUBSCRIBE = 'subscribe'
  UNSUBSCRIBE = 'unsubscribe'


class KrakenWsPrivateClient(object):
  def __init__(
      self,
      key_file,
      *,
      ioloop=None,
      notification_callback=None,
      close_callback=None,
      on_ready=None,
      timeout_sec=10,
      api_url=KRAKEN_FUTURES_WS_API_URL,
  ):

    self._ioloop = ioloop or IOLoop.current()
    self._api_url = api_url
    self._ws = None
    self._ready = False
    self._timeout = timeout_sec * 1e9
    self._auth = KrakenAsyncAuth(key_file, use_extra_keys=False)
    self._ws_token = None

    self._notification_callback = notification_callback or (lambda x: None)
    self._close_callback = close_callback or (lambda: None)
    self._on_ready = on_ready or (lambda: None)
    self._scheduler = Scheduler()

    self._last_received_message_timestamp = None
    self._pending_future = None
    self._subscribed_channels = []

    self._stopped = False
    self._init_params = (ioloop, notification_callback, close_callback, timeout_sec)
    self._logger = logging.getLogger(__name__)

  @property
  def last_received_message_timestamp(self):
    return self._last_received_message_timestamp

  def is_ready(self):
    return self._ready

  def _check_future_timeout(self):
    timestamp = get_timestamp()
    if self._pending_future is None:
      return

    if timestamp - self._pending_future.timestamp > self._timeout:
      self._pending_future.set_exception(
          TimeoutError('Kraken subscription timeout!  %s' % self._pending_future.msg))
      self._pending_future = None

  def _check_connection(self):
    if self._last_received_message_timestamp is None:
      return

    timestamp = get_timestamp()
    diff_time = timestamp - self._last_received_message_timestamp
    if diff_time > self._timeout:
      self._logger.warning('No message received for %s seconds', diff_time / 1e9)
      self.stop()

  @gen.coroutine
  def _connect(self):
    # noinspection PyBroadException
    try:
      self._logger.info('%s try to connect to the websocket.', type(self).__name__)
      self._ws = yield websocket_connect(self._api_url)
      self._ioloop.add_callback(self._run)
      response = yield self.get_ws_token()
      if "result" in response.json():
        self._ws_token = response.json()["result"]["token"]
        self._ready = True
        self._on_ready()
        self._logger.info('Kraken private client websocket connected.')
      else:
        self._logger.error('%s connect to websocket failed! %s',
                           type(self).__name__,
                           json.dumps(response.json()))
    except Exception:
      self._logger.exception('%s connect to websocket failed!', type(self).__name__)
      self.stop()

  @gen.coroutine
  def _run(self):
    while not self._stopped:
      msg = None
      # noinspection PyBroadException
      try:
        msg = yield self._ws.read_message()
      except Exception:
        self._logger.exception('Error reading ws message')
        self.stop()

      if msg is None:
        self._logger.error('Websocket is closed!')
        self.stop()
        return

      if len(msg) == 0:
        self._logger.info('Empty message received!')
        continue

      # noinspection PyBroadException
      try:
        self.on_message(msg)
      except Exception:
        self._logger.exception('Exception raised when process message!')

  def _set_future(self, data):
    if self._pending_future:
      self._pending_future.set_result(data)
      self._pending_future = None
    else:
      self._logger.error('Unexpected msg coming! %s', pretty_json(data))
      self.stop()

  def on_message(self, raw_msg: bytes):
    timestamp = get_timestamp()
    self._last_received_message_timestamp = timestamp
    if self._stopped:
      return

    try:
      data = json.loads(raw_msg)
    except ValueError:
      self._logger.exception('Invalid json string!, %s', raw_msg)
      return

    # data receive
    if isinstance(data, list):
      self._notification_callback(data)
      return

    try:
      event = EventType(data.get('event'))
    except ValueError:
      self._logger.error('Unknown event type: %s', data)
      self.stop()
      return
    except Exception as e:
      self._logger.exception('on_message get event failed! %s', str(e))

    if event == EventType.HEARTBEAT:
      pass
    elif event == EventType.SYSTEM_STATUS:
      formatted = pretty_json(data)
      self._logger.info('systemStatus message received! (future=%s)\n%s',
                        self._pending_future,
                        formatted)
    elif event == EventType.SUBSCRIPTION_STATUS:
      formatted = pretty_json(data)
      if data['status'] == "subscribed":
        self._logger.info('channel subscribed!\n%s', formatted)
        channel = ChannelType(data['channelName'])
        self._subscribed_channels.append(channel)
        self._set_future(data)
      elif data['status'] == "unsubscribed":
        self._logger.info('channel unsubscribed!\n%s', formatted)
        channel = ChannelType(data['channelName'])
        self._subscribed_channels.remove(channel)
        self._set_future(data)
      elif data['status'] == "error":
        formatted = pretty_json(data)
        self._logger.error('error msg received!\n%s', formatted)
        self.stop()
    else:
      formatted = pretty_json(data)
      self._logger.error('Unknown event type!\n%s', formatted)

  def _write_message(self, msg):
    try:
      self._ws.write_message(json.dumps(msg))
    except WebSocketClosedError as e:
      self._logger.error('Write message error, ws is close!')
      self._pending_future.set_exception(e)
      self._pending_future = None

  def send_request(self, msg):
    if self._ws is None:
      self._logger.error('websocket not connected yet! %s', msg)
      return

    if self._pending_future is not None:
      self._logger.error('Pending future is not resolved yet! %s', msg)
      return

    fut = Future()
    fut.timestamp = get_timestamp()
    fut.msg = msg
    self._logger.info('send request: %s', msg)
    self._ioloop.add_callback(self._write_message, msg)
    self._pending_future = fut
    return fut

  def get_ws_token(self):
    return self._auth.query("GetWebSocketsToken")

  def subscribe(self, channel: ChannelType):
    self._logger.info('Subscribe topic (%s)\n', channel)
    if self._ws_token is None:
      self._logger.error('ws_token is not solved!')
      return

    msg = {
        'event': 'subscribe', 'subscription': {
            'name': channel.value,
            'token': self._ws_token,
        }
    }
    return self.send_request(msg)

  def unsubscribe(self, channel: ChannelType):
    self._logger.info('Unsubscribe topic (%s)\n', channel)
    if self._ws_token is None:
      self._logger.error('ws_token is not solved!')
      return

    msg = {
        'event': 'unsubscribe', 'subscription': {
            'name': channel.value,
            'token': self._ws_token,
        }
    }
    return self.send_request(msg)

  def start(self):
    self._ioloop.add_callback(self._connect)
    self._scheduler.add_periodic_task(1, self._check_future_timeout)
    self._scheduler.add_periodic_task(10, self._check_connection)
    self._scheduler.start()

  def stop(self):
    if self._stopped:
      return
    self._scheduler.stop()
    self._ready = False
    self._ioloop.add_callback(self._async_stop)
    self._logger.info('Kraken ws private client (%s) stopped!', id(self))

  @gen.coroutine
  def _async_stop(self):
    for channel in self._subscribed_channels:
      response = yield self.unsubscribe(channel)
      self._logger.info('Topic (%s) unsubscribed!\n%s', channel, response)

    self._stopped = True
    if self._ws is not None:
      self._ws.close()
      self._ws = None
    yield gen.sleep(10)
    self._close_callback()
